/*
 * Copyright 2014 Google Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.inferred.freebuilder.processor.util;

import com.google.common.annotations.VisibleForTesting;
import com.google.googlejavaformat.java.Formatter;

import org.inferred.freebuilder.processor.util.feature.Feature;
import org.inferred.freebuilder.processor.util.feature.FeatureSet;
import org.inferred.freebuilder.processor.util.feature.FeatureType;

import java.util.Collection;

import javax.annotation.processing.ProcessingEnvironment;
import javax.lang.model.element.PackageElement;
import javax.lang.model.element.TypeElement;
import javax.lang.model.util.ElementFilter;

/** {@code SourceBuilder} which also handles package declaration and imports. */
public class CompilationUnitBuilder implements SourceBuilder {

  private final ImportManager importManager;
  private final SourceBuilder source;
  private final QualifiedName classToWrite;

  /**
   * Returns a {@link CompilationUnitBuilder} for {@code classToWrite} using {@code features}. The
   * file preamble (package and imports) will be generated automatically, and {@code env} will be
   * inspected for potential import collisions.
   */
  public CompilationUnitBuilder(
      ProcessingEnvironment env,
      QualifiedName classToWrite,
      Collection<QualifiedName> nestedClasses,
      FeatureSet features) {
    this.classToWrite = classToWrite;
    // Write the source code into an intermediate SourceStringBuilder, as the imports need to be
    // written first, but aren't known yet.
    ImportManager.Builder importManagerBuilder = new ImportManager.Builder();
    importManagerBuilder.addImplicitImport(classToWrite);
    PackageElement pkg = env.getElementUtils().getPackageElement(classToWrite.getPackage());
    for (TypeElement sibling : ElementFilter.typesIn(pkg.getEnclosedElements())) {
      importManagerBuilder.addImplicitImport(QualifiedName.of(sibling));
    }
    for (QualifiedName nestedClass : nestedClasses) {
      importManagerBuilder.addImplicitImport(nestedClass);
    }
    importManager = importManagerBuilder.build();
    source = new SourceStringBuilder(importManager, features);
  }

  @Override
  public CompilationUnitBuilder add(String fmt, Object... args) {
    source.add(fmt, args);
    return this;
  }

  @Override
  public SourceBuilder add(Excerpt excerpt) {
    source.add(excerpt);
    return this;
  }

  @Override
  public CompilationUnitBuilder addLine(String fmt, Object... args) {
    source.addLine(fmt, args);
    return this;
  }

  @Override
  public SourceStringBuilder subBuilder() {
    return source.subBuilder();
  }

  @Override
  public <T extends Feature<T>> T feature(FeatureType<T> feature) {
    return source.feature(feature);
  }

  @Override
  public String toString() {
    StringBuilder unit = new StringBuilder();
    unit.append("// Autogenerated code. Do not modify.\n")
        .append("package ").append(classToWrite.getPackage()).append(";\n")
        .append("\n");
    if (!importManager.getClassImports().isEmpty()) {
      for (String classImport : importManager.getClassImports()) {
        unit.append("import ").append(classImport).append(";\n");
      }
      unit.append("\n");
    }
    unit.append(formatSource(source.toString()));
    return unit.toString();
  }

  @VisibleForTesting
  public static String formatSource(String source) {
    try {
      return new Formatter().formatSource(source);
    } catch (UnsupportedClassVersionError e) {
      // Formatter requires Java 7+; do no formatting in Java 6.
      return source;
    } catch (Exception e) {
      StringBuilder message = new StringBuilder()
          .append("Formatter failed:\n")
          .append(e.getMessage())
          .append("\nGenerated source:");
      int lineNo = 0;
      for (String line : source.split("\n")) {
        message
            .append("\n")
            .append(++lineNo)
            .append(": ")
            .append(line);
      }
      throw new RuntimeException(message.toString());
    }
  }
}
